package com.example.daily.ui.foodrecognition

import android.content.Context
import android.graphics.Bitmap
import android.net.Uri
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import com.example.daily.data.dao.FoodRecordDao
import com.example.daily.data.UserPreferences
import dagger.hilt.android.lifecycle.HiltViewModel
import dagger.hilt.android.qualifiers.ApplicationContext
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.launch
import javax.inject.Inject
import android.graphics.ImageDecoder
import android.os.Build
import android.provider.MediaStore
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.support.common.FileUtil
import org.tensorflow.lite.support.image.ImageProcessor
import org.tensorflow.lite.support.image.TensorImage
import org.tensorflow.lite.support.image.ops.ResizeOp
import org.tensorflow.lite.support.tensorbuffer.TensorBuffer
import java.nio.ByteBuffer
import java.nio.ByteOrder

data class FoodRecognitionUiState(
    val isLoading: Boolean = false,
    val error: String? = null,
    val selectedImageUri: Uri? = null,
    val recognitionResults: List<RecognitionResult> = emptyList()
)

data class RecognitionResult(
    val label: String,
    val confidence: Float,
    val calories: Int = 0
)

@HiltViewModel
class FoodRecognitionViewModel @Inject constructor(
    @ApplicationContext private val context: Context,
    private val foodRecordDao: FoodRecordDao,
    private val userPreferences: UserPreferences
) : ViewModel() {

    private val _uiState = MutableStateFlow(FoodRecognitionUiState())
    val uiState: StateFlow<FoodRecognitionUiState> = _uiState

    private var interpreter: Interpreter? = null
    private val labels = listOf("nata", "pasteis_bacalhau", "pizza", "tripas_moda_porto", "waffles")
    
    // 模型配置
    private val imageWidth = 224
    private val imageHeight = 224
    private val numChannels = 3
    private val batchSize = 1
    private val pixelSize = 4 // Float32 = 4 bytes

    private val imageProcessor = ImageProcessor.Builder()
        .add(ResizeOp(imageHeight, imageWidth, ResizeOp.ResizeMethod.BILINEAR))
        .build()

    init {
        initializeInterpreter()
    }

    private fun initializeInterpreter() {
        try {
            val modelFile = FileUtil.loadMappedFile(context, "food_model.tflite")
            val options = Interpreter.Options().apply {
                setNumThreads(4)
            }
            interpreter = Interpreter(modelFile, options)
        } catch (e: Exception) {
            _uiState.value = _uiState.value.copy(
                error = "模型加载失败: ${e.message}"
            )
        }
    }

    fun setSelectedImage(uri: Uri) {
        _uiState.value = _uiState.value.copy(
            selectedImageUri = uri,
            isLoading = true,
            error = null
        )
        recognizeFood(uri)
    }

    private fun recognizeFood(uri: Uri) {
        viewModelScope.launch {
            try {
                val bitmap = loadBitmap(uri)
                val results = withContext(Dispatchers.Default) {
                    // 准备输入数据
                    val inputBuffer = ByteBuffer.allocateDirect(batchSize * imageHeight * imageWidth * numChannels * pixelSize)
                    inputBuffer.order(ByteOrder.nativeOrder())
                    
                    // 将图像转换为浮点数据
                    val intValues = IntArray(imageWidth * imageHeight)
                    bitmap.getPixels(intValues, 0, imageWidth, 0, 0, imageWidth, imageHeight)
                    
                    inputBuffer.rewind()
                    for (pixelValue in intValues) {
                        // 归一化到[-1, 1]范围
                        inputBuffer.putFloat(((pixelValue shr 16 and 0xFF) - 127.5f) / 127.5f) // R
                        inputBuffer.putFloat(((pixelValue shr 8 and 0xFF) - 127.5f) / 127.5f)  // G
                        inputBuffer.putFloat(((pixelValue and 0xFF) - 127.5f) / 127.5f)        // B
                    }
                    
                    // 准备输出数据
                    val outputBuffer = TensorBuffer.createFixedSize(intArrayOf(1, labels.size), org.tensorflow.lite.DataType.FLOAT32)
                    
                    // 运行推理
                    interpreter?.run(inputBuffer, outputBuffer.buffer.rewind())
                    
                    // 处理结果
                    val probabilities = outputBuffer.floatArray
                    probabilities.mapIndexed { index, probability ->
                        RecognitionResult(
                            label = labels[index],
                            confidence = probability,
                            calories = getCaloriesForFood(labels[index])
                        )
                    }.sortedByDescending { it.confidence }
                        .take(3)
                        .filter { it.confidence > 0.3f }
                }
                
                _uiState.value = _uiState.value.copy(
                    recognitionResults = results,
                    isLoading = false
                )
            } catch (e: Exception) {
                _uiState.value = _uiState.value.copy(
                    error = "识别失败: ${e.message}",
                    isLoading = false
                )
            }
        }
    }

    private suspend fun loadBitmap(uri: Uri): Bitmap {
        return withContext(Dispatchers.IO) {
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.P) {
                val source = ImageDecoder.createSource(context.contentResolver, uri)
                ImageDecoder.decodeBitmap(source) { decoder, _, _ ->
                    decoder.isMutableRequired = true
                    decoder.setTargetSize(imageWidth, imageHeight)
                }
            } else {
                @Suppress("DEPRECATION")
                MediaStore.Images.Media.getBitmap(context.contentResolver, uri).let { originalBitmap ->
                    Bitmap.createScaledBitmap(originalBitmap, imageWidth, imageHeight, true)
                }
            }
        }
    }

    private fun getCaloriesForFood(foodName: String): Int {
        return when (foodName) {
            "nata" -> 300          // 葡式蛋挞
            "pasteis_bacalhau" -> 180  // 鳕鱼饼
            "pizza" -> 266         // 披萨（每片）
            "tripas_moda_porto" -> 350 // 波尔图风格牛肚
            "waffles" -> 250       // 华夫饼
            else -> 200
        }
    }

    override fun onCleared() {
        super.onCleared()
        interpreter?.close()
    }
} 